/*
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.facebook.presto.mongodb.optimization;

import com.facebook.presto.mongodb.MongoTableHandle;
import com.facebook.presto.mongodb.MongoTableLayoutHandle;
import com.facebook.presto.spi.ColumnHandle;
import com.facebook.presto.spi.TableHandle;
import com.facebook.presto.spi.adaptive.CatalogEngine;
import com.facebook.presto.spi.adaptive.ConnectorAdaptiveOptimizer;
import com.facebook.presto.spi.adaptive.OptimizerType;
import com.facebook.presto.spi.adaptive.metadata.presto.AdaptivePrestoObject;
import com.facebook.presto.spi.plan.FilterNode;
import com.facebook.presto.spi.plan.PlanNode;
import com.facebook.presto.spi.plan.PlanNodeIdAllocator;
import com.facebook.presto.spi.plan.TableScanNode;
import com.facebook.presto.spi.predicate.Domain;
import com.facebook.presto.spi.predicate.TupleDomain;
import com.facebook.presto.spi.relation.VariableReferenceExpression;

import java.util.HashMap;
import java.util.Map;
import java.util.Optional;

public class MongoUnionIndexPushdown
        implements ConnectorAdaptiveOptimizer
{
    @Override
    public CatalogEngine getConnectorCatalog()
    {
        return CatalogEngine.MONGODB;
    }

    @Override
    public OptimizerType getOptimizerType()
    {
        return OptimizerType.UNIONINDEX;
    }

    @Override
    public PlanNode adaptiveOptimize(
            PlanNode node,
            TableHandle oldTableHandle,
            PlanNodeIdAllocator idAllocator,
            AdaptivePrestoObject adaptiveObject)
    {
        if (node instanceof TableScanNode) {

            TableScanNode oldTableScanNode = (TableScanNode) node;
            TableScanNode newTableScanNode = replaceTableScanNode(oldTableScanNode, oldTableHandle, idAllocator, adaptiveObject);
            if (oldTableScanNode == newTableScanNode) {
                return node;
            }
            return newTableScanNode;
        }

        if (node instanceof FilterNode) {
            TableScanNode oldTableScanNode = (TableScanNode) ((FilterNode) node).getSource();
            TableScanNode newTableScanNode = replaceTableScanNode(oldTableScanNode, oldTableHandle, idAllocator, adaptiveObject);
            if (oldTableScanNode == newTableScanNode) {
                return node;
            }
            // create newFilterNode and return
            return new FilterNode(idAllocator.getNextId(), newTableScanNode, ((FilterNode) node).getPredicate());
        }
        return node;
    }

    private TableScanNode replaceTableScanNode(
            TableScanNode oldTableScanNode,
            TableHandle oldTableHandle,
            PlanNodeIdAllocator idAllocator,
            AdaptivePrestoObject adaptiveObject)
    {
        // get adaptive object.
        TableHandle storedTableHandle = adaptiveObject.getTableHandle();
        Map<String, ColumnHandle> storedColumnHandle = adaptiveObject.getColumnHandles();

        //TODO: tupleDomain is null?
        if (!oldTableHandle.getLayout().isPresent()) {
            return oldTableScanNode;
        }

        // new tupleDomain and replace the old.

        // 1. getOldTupleDomain
        // 2. oldDomains: <oldColumnHandle, domain>
        // ==> newDomains: <MongodbColumnHandle, domain>
        // 3. create newTupleDomain
        TupleDomain<ColumnHandle> oldTupleDomain = oldTableHandle.getLayout().get().getTupleDomain();
        Map<ColumnHandle, Domain> newDomains = new HashMap<>();

        Map<ColumnHandle, Domain> oldDomains = oldTupleDomain.getDomains().get();

        oldDomains.forEach((oldHandle, oldDomain) -> newDomains.put(storedColumnHandle.get(oldHandle.getColumnName()), oldDomain));

        TupleDomain<ColumnHandle> newTupleDomain = TupleDomain.withColumnDomains(newDomains);

        // create newTableLayoutHandle
        MongoTableLayoutHandle newTableLayoutHandle =
                new MongoTableLayoutHandle((MongoTableHandle) storedTableHandle.getConnectorHandle(), newTupleDomain, "addition");

        // create newTableHandle
        TableHandle newTableHandle = new TableHandle(
                storedTableHandle.getConnectorId(),
                storedTableHandle.getConnectorHandle(),
                storedTableHandle.getTransaction(),
                Optional.of(newTableLayoutHandle));

        // create newAssignments
        Map<VariableReferenceExpression, ColumnHandle> newAssignments = new HashMap<>();

        oldTableScanNode.getOutputVariables()
                .forEach(expression -> newAssignments.put(expression, storedColumnHandle.get(expression.getName())));

        // create newTableScanNode
        TableScanNode newTableScanNode = new TableScanNode(
                idAllocator.getNextId(),
                newTableHandle,
                oldTableScanNode.getOutputVariables(),
                newAssignments,
                //TODO: constraint needs to change columnHandle.
                oldTableScanNode.getCurrentConstraint(),
                oldTableScanNode.getEnforcedConstraint()
        );
        return newTableScanNode;
    }
}
